// SPDX-License-Identifier: Mulan PSL v2
/*
 * Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This software is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *         http://license.coscl.org.cn/MulanPSL2
 *
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

#![allow(unused)]

use tracing::debug;
use xgpu_common::api_name::ApiFuncName;
use xgpu_common::ipc::message::Request;

use crate::{agent, fault_guard::journal, thread};

const RECOVERY_HEALTHY_SKIP_APIS: [ApiFuncName; 7] = [
    ApiFuncName::FuncCudamalloc,
    ApiFuncName::FuncCudamemcpyasync,
    ApiFuncName::FuncCudamemsetasync,
    ApiFuncName::FuncNcclcomminitrankconfigrecovery,
    ApiFuncName::FuncCumoduleloaddata,
    ApiFuncName::FuncCumodulegetfunction,
    ApiFuncName::FuncCudafree,
];

pub fn recovery_healthy() -> anyhow::Result<()> {
    debug!("[Recovery] Start recovery (healthy rank)");

    while let Some(mut req_bytes) = journal::load_next_request()? {
        let mut req: Request = journal::deserialize_request(&mut req_bytes)?;
        let req_id = req.request_id();
        let method_id = req.method_id();
        let api = ApiFuncName::try_from(method_id as i32)?;
        debug!(
            "[Recovery] Stage 1, replaying api: request_id={}, method_id={}",
            req_id, method_id
        );

        if !RECOVERY_HEALTHY_SKIP_APIS.contains(&api) {
            thread::context().agent().replay_api(&mut req)?;
        }
    }

    while let Some(mut req_bytes) = journal::fetch_minibatch_request()? {
        let mut req: Request = journal::deserialize_request(&mut req_bytes)?;
        let req_id = req.request_id();
        let method_id = req.method_id();
        debug!(
            "[Recovery] Stage 2, replaying api: request_id={}, method_id={}",
            req_id, method_id
        );

        if method_id != ApiFuncName::FuncCudamalloc as u64 {
            thread::context().agent().replay_api(&mut req)?;
        }
    }

    Ok(())
}

pub fn recovery_failed() -> anyhow::Result<()> {
    debug!("[Recovery] Start recovery (failed rank)");

    while let Some(mut req_bytes) = journal::load_next_request()? {
        let mut req: Request = journal::deserialize_request(&mut req_bytes)?;
        let req_id = req.request_id();
        let method_id = req.method_id();
        debug!(
            "[Recovery] Stage 1, replaying api: request_id={}, method_id={}",
            req_id, method_id
        );

        thread::context().agent().replay_api(&mut req)?;
    }

    while let Some(mut req_bytes) = journal::fetch_minibatch_request()? {
        let mut req: Request = journal::deserialize_request(&mut req_bytes)?;
        let req_id = req.request_id();
        let method_id = req.method_id();
        debug!(
            "[Recovery] Stage 2, replaying api: request_id={}, method_id={}",
            req_id, method_id
        );

        thread::context().agent().replay_api(&mut req)?;
    }

    Ok(())
}
